'''
Created on 16 mars 2011

@author: wp803469

    Manipulation FFT: question 4.4
    

'''

import matplotlib.pyplot as plt 
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import numpy as np
import Image

def main():
    
    im = Image.open('train.png').convert('I') # pas besoin de la couche alpha
    im = im.rotate(90)
    SZ = im.size[0]
    print im
    print "debug de l'image:",im.format, im.size, im.mode
    
    plt.figure(1)
    plt.clf()
    plt.imshow(im, cmap=plt.cm.gray_r, interpolation="Nearest",origin="lower")
    #plt.show()
    
    # fft
    F1 = np.fft.fft2(im)

    
    F = np.ndarray(shape=(len(F1), len(F1[0])))
    i=0
    j=0

    for c1 in F1:
        j=0
        for c2 in c1:
            moduleDeC = np.abs(c2)
            #print "module de c2: ",moduleDeC
            
            if moduleDeC < 5000:
                print "choped !"
                F[j][i] = np.complex128(0 + 0j)
            else :
                F[j][i] = c2
                
            #F[i][j] = moduleDeC
            j = j+1
        i = i+1

    F2 = np.fft.fftshift(F)
    
    # Calculate a 2D magnitude spectrum
    '''
    psd2D =np.abs(F2)
    fig = plt.figure(2)
    plt.clf()
    X = range(0,SZ)
    Y = range(0,SZ)
    X, Y = np.meshgrid(X, Y)
    ax = Axes3D(fig)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.plot_surface(X, Y, psd2D, rstride=1, cstride=1, cmap=cm.jet)
    '''
    
    # Inverse fft
    imm = np.fft.ifft2(F)
    plt.figure(2)
    plt.clf()
    plt.imshow(imm.real, cmap=cm.gray)

    plt.show()

    




if __name__ == '__main__':
    main()